from BaseBlocks import ImageCapture,ImageProcessor
import cv2 as cv
import numpy as np

class ImageCapture_OpenCV(ImageCapture):
    def __init__(self,
                deviceID,
                imageBufferSize=5):
        '''这里将摄像头V4L2摄像头对象例化
        '''
        super().__init__(imageBufferSize)
        self.videoCaptureDevice = cv.VideoCapture(deviceID,cv.CAP_V4L2)

    def captureTask(self):
        return self.videoCaptureDevice.read()

class ImageProcessor_ConvertToGray(ImageProcessor):
    def __init__(self, name):
        super().__init__(name)
    def process(self, inputImage, args):
        image = cv.cvtColor(inputImage,cv.COLOR_BGR2GRAY)
        return image

class ImageProcessor_OtsuThreshould(ImageProcessor):
    def __init__(self, name):
        super().__init__(name)
    def process(self, inputImage, args):
        _,image = cv.threshold(inputImage,0,255,cv.THRESH_OTSU)
        return image
    
class ImageProcessor_HSVRange(ImageProcessor):
    def __init__(self, name):
        super().__init__(name)
        self.requiredArgs = [
            '%s.lowerH'%self.name,
            '%s.upperH'%self.name,
            '%s.lowerS'%self.name,
            '%s.upperS'%self.name,
            '%s.lowerV'%self.name,
            '%s.upperV'%self.name
            ]
    def process(self, inputImage, args):
        # 这些是必须设置的参数
        lowerH = args[self.requiredArgs[0]]
        upperH = args[self.requiredArgs[1]]
        lowerS = args[self.requiredArgs[2]]
        upperS = args[self.requiredArgs[3]]
        lowerV = args[self.requiredArgs[4]]
        upperV = args[self.requiredArgs[5]]

        image = cv.cvtColor(inputImage,cv.COLOR_BGR2HSV)
        lowerRange = np.array([lowerH,lowerS,lowerV])
        upperRange = np.array([upperH,upperS,upperV])
        imageInRange = cv.inRange(image,lowerRange,upperRange)
        
        return imageInRange

class ImageAnalyzer_Centroid(ImageProcessor):
    def __init__(self, name):
        super().__init__(name)
    def process(self, inputImage, args):
        momentDict = cv.moments(inputImage,True)
        if momentDict['m00'] != 0:
            return momentDict['m01']/momentDict['m00'],momentDict['m10']/momentDict['m00']
        else:
            return None